import numpy as np
from 参数输入 import define_material_and_geometry
from 节点和单元 import define_nodes_and_elements
from 局部单元刚度矩阵 import beam_element_stiffness_curved_3d
from 单刚汇入总刚 import assemble_global_stiffness
from 边界条件 import apply_boundary_conditions
from 求解节点位移 import solve_displacements
from 可视化三维模型 import plot_helix

def main():
    # 1. 定义材料和几何属性
    E, Ixx, Iyy, J, A, R, theta, h = define_material_and_geometry()

    # 2. 定义节点和单元
    num_nodes = 20 # 节点数量
    nodes, elements, nodes_angle = define_nodes_and_elements(num_nodes, R, theta, h)
    num_elements = num_nodes - 1
    dof_per_node = 6  # 每个节点的自由度数量

    # 3. 初始化全局刚度矩阵和力向量
    K = np.zeros((num_nodes * dof_per_node, num_nodes * dof_per_node))  # 全局刚度矩阵
    F = np.zeros(num_nodes * dof_per_node)  # 力向量

    # 4. 计算单元刚度矩阵并汇总到全局刚度矩阵
    for element in elements:
        node1, node2 = element
        L_elem = R * (nodes_angle[node2] - nodes_angle[node1])  # 弧长 = 半径 * 角度差
        k_local = beam_element_stiffness_curved_3d(E, Ixx, Iyy, J, A, L_elem, R)
        #assemble_global_stiffness(K, k_local, element, dof_per_node)
        assemble_global_stiffness(K, k_local, element, dof_per_node, nodes)

    # 5. 应用边界条件
    boundary_conditions = [{'dof': i, 'value': 0.0} for i in range(dof_per_node)]  # 固定左端 (节点0的6个自由度)
    apply_boundary_conditions(K, F, boundary_conditions)

    # 6. 定义力向量（右端施加一个集中力）
    F[-5] = 1000  # 节点最后一个节点的Y方向力 (集中力 1000N)

    # 7. 求解位移
    displacements = solve_displacements(K, F)

    # 输出节点位移
    print("节点位移：")
    for i in range(num_nodes):
        disp = displacements[i * dof_per_node:(i + 1) * dof_per_node]
        print(f"节点 {i} 的位移: {disp[:3]}，旋转: {disp[3:]}")

    # 8. 可视化三维曲梁螺旋模型
    plot_helix(nodes)

if __name__ == "__main__":
    main()
